import csv
import collections
import sys
from datetime import datetime
from tqdm import tqdm
from time import time
from util import is_app, def_user

SRC_PATH = "./data/sub_train_100k.csv"
DST_TRAIN_PATH = "./data/new_train_sample.csv"
DST_VALID_PATH = "./data/new_valid_sample.csv"

TRAIN_APP_PATH = "./data/new_train_app.csv"
VALID_APP_PATH = "./data/new_valid_app.csv"
TRAIN_SITE_PATH = "./data/new_train_site.csv"
VALID_SITE_PATH = "./data/new_valid_site.csv"

FIELDS = [
    'id', 'click', 'hour','C1','banner_pos','site_id','site_domain',
    'site_category','app_id', 'app_domain','app_category', 'device_id',
    'device_ip', 'device_model', 'device_type', 'device_conn_type', 'C14',
    'C15','C16','C17','C18','C19','C20','C21'
]

NEW_FIELDS = FIELDS+['user','pub_id','pub_domain','pub_category','week_day','hour_day','device_id_count','device_ip_count','user_count','user_hour_count','pub_user_cnt']

def split_train_valid(src_path, dst_train_path, dst_valid_path):
    """
    将数据切分为训练集与验证集
    """
    with open(src_path, 'r') as csv_file:
        reader = csv.DictReader(csv_file, FIELDS)
        writer_train = csv.DictWriter(open(dst_train_path, 'w', newline=''), FIELDS)
        writer_valid = csv.DictWriter(open(dst_valid_path, 'w', newline=''), FIELDS)

        writer_train.writeheader()
        writer_valid.writeheader()
        cnt = 0
        for i, row in tqdm(enumerate(reader)):

            if i == 0:
                continue
            new_row = {}
            for field in FIELDS:
                new_row[field] = row[field]
            hour = datetime.strptime(row['hour'], "%y%m%d%H")
            if hour.day <= 28:
                writer_train.writerow(new_row)
            elif hour.day > 28:
                writer_valid.writerow(new_row)

id_cnt = collections.defaultdict(int)
ip_cnt = collections.defaultdict(int)

user_cnt = collections.defaultdict(int)
user_hour_cnt = collections.defaultdict(int)

site_id_users_set = collections.defaultdict(set)
app_id_users_set = collections.defaultdict(set)
start = time()


def scan(path):
    """
    数据量统计
    """
    for i, row in enumerate(csv.DictReader(open(path)), start=1):
        if i % 1000000 == 0:
            sys.stderr.write('{0:6.0f}    {1}m\n'.format(time() - start, int(i / 1000000)))

        user = def_user(row)
        id_cnt[row['device_id']] += 1
        ip_cnt[row['device_ip']] += 1
        user_cnt[user] += 1
        user_hour_cnt[str(user) + '-' + row['hour']] += 1

        if is_app(row):
            app_id_users_set[row['app_id']].add(user)
        else:
            site_id_users_set[row['site_id']].add(user)

def gen_data(src_path, dst_app_path, dst_site_path, is_train=True):
    """
    将数据按广告来源切分为 app 类型或 site 类型，并新增特征
    """
    reader = csv.DictReader(open(src_path))
    writer_app = csv.DictWriter(open(dst_app_path, 'w', newline=''), NEW_FIELDS)
    writer_site = csv.DictWriter(open(dst_site_path, 'w', newline=''), NEW_FIELDS)
    writer_app.writeheader()
    writer_site.writeheader()

    for i, row in tqdm(enumerate(reader, start=1)):
        new_row = {}
        if is_train:
            for field in FIELDS:
                new_row[field] = row[field]
        else:
            for field in FIELDS:
                if field == 'click':
                    continue
                new_row[field] = row[field]
        new_row['device_id_count'] = id_cnt[row['device_id']]
        new_row['device_ip_count'] = ip_cnt[row['device_ip']]
        user, hour = def_user(row), row['hour']
        new_row['user'] = user
        new_row['user_count'] = user_cnt[user]
        new_row['user_hour_count'] = user_hour_cnt[str(user) + '-' + hour]

        row_date = datetime.strptime(str(hour), "%y%m%d%H")
        new_row['week_day'] = row_date.weekday()
        new_row['hour_day'] = row_date.hour

        if is_app(row):
            new_row['pub_id'] = row['app_id']
            new_row['pub_domain'] = row['app_domain']
            new_row['pub_category'] = row['app_category']
            new_row['pub_user_cnt'] = len(app_id_users_set[row['app_id']])
            writer_app.writerow(new_row)
        else:
            new_row['pub_id'] = row['site_id']
            new_row['pub_domain'] = row['site_domain']
            new_row['pub_category'] = row['site_category']
            new_row['pub_user_cnt'] = len(site_id_users_set[row['site_id']])
            writer_site.writerow(new_row)

if __name__ == '__main__':
    # 训练集与验证集切分
    split_train_valid(SRC_PATH, DST_TRAIN_PATH, DST_VALID_PATH)

    # Counting Features 统计量
    scan(DST_TRAIN_PATH)
    scan(DST_VALID_PATH)

    # 按 App 与 Site 切分数据并新增特征
    gen_data(DST_TRAIN_PATH, TRAIN_APP_PATH, TRAIN_SITE_PATH)
    gen_data(DST_VALID_PATH, VALID_APP_PATH, VALID_SITE_PATH)